'''
 # @ Author: Justin
 # @ Create Time: 2025-02-11 19:48:25
 # @ Modified by: Justin
 # @ Modified time: 2025-02-11 19:48:28
 # @ Description: 与deepseek的聊天记录
 '''

import json
import sqlite3
import time
import os

class Messages:
    def __init__(self, path : str = None, prompt : str = None) :
        '''
        Args:
            path(str, optional): 聊天记录文件路径. 默认为None.
                                如果为None，则不在文件中存储聊天记录。
                                如果为文件路径不存在，则创建文件。
            prompt(str, optional): 系统提示. 默认为None.
        '''

        '''
        数据库message表结构：
        字段         类型               说明
        id           INTEGER PRIMARY    主键
        time         TEXT               时间
        role         TEXT               角色
        message      TEXT               聊天内容
        '''

        '''
        messages列表结构：
        [
            {
                "id"
                "time"
                "role"
                "message"
            },
            {
            ...
            },
            ...
        ]
        '''
        self.messages = []
        self.prompt = prompt

        # 若path不为None，则创建数据库文件
        if path is not None:
            self.saveMethod = 'db'  # 保存方式为数据库
            self.path = path

            try:
                # 连接数据库
                conn = sqlite3.connect(f'{self.path}')
                c = conn.cursor()

                # 创建message表
                c.execute('''CREATE TABLE IF NOT EXISTS message
                            (id INTEGER PRIMARY KEY AUTOINCREMENT,
                            time DATETIME,
                            role TEXT,
                            message TEXT)''')
                conn.commit()
                conn.close()
            except sqlite3.Error as e:
                print(f'连接数据库失败: {e}')
            
            # 从数据库加载聊天记录
            self.load_message_from_db()
        else:
            self.saveMethod = 'list'  # 保存方式为列表
            self.path = None
            self.max_id = -1

        # 添加系统提示
        if self.prompt is not None and self.max_id != -1:
            print(f'添加系统提示\n"{prompt}"\n错误：不能在已开启的对话上添加系统提示')
            self.prompt = None
        elif self.prompt is not None:
            self.add_message(time=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), role='system', message=self.prompt)

    def load_message_from_db(self):
        '''
        从数据库中加载聊天记录
        '''

        try:
            # 连接数据库
            conn = sqlite3.connect(f'{self.path}')
            c = conn.cursor()

            # 查询message表
            c.execute('SELECT * FROM message')
            rows = c.fetchall()

            self.messages = [{"id": row[0], "time": row[1], "role": row[2], "message": row[3]} for row in rows]
            self.max_id = len(self.messages) - 1

            # 加载系统提示
            c.execute('SELECT message FROM message WHERE role = "System"')
            prompt = c.fetchone()

            if prompt is not None:
                self.prompt = prompt[0]

            conn.close()

        except sqlite3.Error as e:
            print(f'查询数据库失败: {e}')
            return

    def save_message_to_db(self):
        '''
        保存聊天记录到数据库
        '''
        if self.saveMethod == 'list':
            return
        
        try:
            # 连接数据库
            conn = sqlite3.connect(f'{self.path}')
            c = conn.cursor()
            
            # 获得最大id
            c.execute('SELECT MAX(id) FROM message')
            max_id = c.fetchone()[0]
            if max_id is None:
                max_id = -1

            # 插入message表
            for message in self.messages:
                if message['id'] > max_id:
                    c.execute("INSERT INTO message (time, role, message) VALUES (?,?,?)", (message['time'], message['role'], message['message']))

            conn.commit()
            conn.close()

        except sqlite3.Error as e:
            print(f'插入数据库失败: {e}')
            return



    def get_messages(self) -> list:
        '''
        获取聊天记录
        Returns:
            list: 聊天记录列表
        '''
        
        return self.messages
    
    def add_message(self, time : str, role : str, message : str) -> int:
        '''
        添加聊天记录
        Args:
            time(str): 时间
            role(str): 角色
            message(str): 聊天内容
        Returns:
            int: 聊天记录id
        '''

        self.messages.append({"id": self.max_id + 1, "time": time, "role": role, "message": message})
        self.max_id += 1

        return self.max_id

    def get_chat_messages(self):
        '''
        获取能够传入deepseek的聊天记录
        Returns:
            list: 聊天记录列表
        '''

        chat_messages = []

        for message in self.messages:
            chat_messages.append({'role': message['role'],'content': message['message']})

        return chat_messages
    
    def delete_self(self) -> None:
        '''
        删除自己（包括数据库文件）
        '''

        if self.saveMethod == 'db':
            try:
                os.remove(self.path)
            except OSError as e:
                print(f'删除数据库文件失败: {e}')

        self.messages = []
        self.max_id = -1
        self.prompt = None

    def __repr__(self):
        return json.dumps(self.messages, indent=4)